import glob
import numpy as np
import cv2
import os


IMAGE_PATH = "D:/vllm/cifar-10-batches-py"
def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


def write_im(im_type, out_folder):
    label_name = ["airplane", "automobile", "bird", "cat", "deer",
                  "dog", "frog", "horse", "ship", "truck"]
    data_batch = "{}/{}".format(IMAGE_PATH, im_type)
    train_list = glob.glob(data_batch)
    save_path = "{}/{}".format(IMAGE_PATH, out_folder)
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    for l in train_list:
        l_dict = unpickle(l)
        for im_idx, im_data in enumerate(l_dict[b'data']):
            im_label = l_dict[b'labels'][im_idx]
            im_name = l_dict[b'filenames'][im_idx]
            im_label_name = label_name[im_label]
            im_data = np.reshape(im_data, [3, 32, 32])
            im_data = np.transpose(im_data, (1, 2, 0))

            #cv2.imshow("im_data", cv2.resize(im_data, (100, 100)))
            #cv2.waitKey(0)
            label_path = "{}/{}".format(save_path, im_label_name)
            if not os.path.exists(label_path):
                os.mkdir(label_path)
            cv2.imwrite("{}/{}".format(label_path, im_name.decode("utf-8")), im_data)


write_im("data_batch_*", "train")
write_im("test_batch", "test")
